import numpy as np
import matplotlib.pyplot as plt
import numpy.random as npr
import matplotlib
matplotlib.use(backend="TkAgg")

npr.seed(0)

# 参数
trials = 20           # 蒙特卡洛模拟次数
p = 0.5                    # 公平硬币
max_m = 50                 # payoff 上限参数 payoff 最多到 2^m

# 模拟首次出现正面的试验次数 N ~ Geometric(p)
N = npr.geometric(p, size=trials)

# -----------------------------
# 图 1: 截断期望随 m 的增长
# -----------------------------
ms = np.arange(1, max_m + 1)
empirical_means = np.empty_like(ms, dtype=float)
for i, m in enumerate(ms):
    '''
    1.如果真实的N<=m,payoff 是 2^N
    2.如果真实的N>m,payoff 是 2^m
    这样实现了“支付上限为2^m的情形”
    '''
    payoffs = 2 ** np.minimum(N, m)  # 若没出现正面则最多支付 2^m
    empirical_means[i] = payoffs.mean()

# 理论结果 E[X] = m + 1
theoretical_means = ms + 1

plt.figure(figsize=(9,5))
plt.plot(ms, empirical_means, label='Empirical mean (cap at 2^m)', marker='o')
plt.plot(ms, theoretical_means, label='Theoretical mean = m+1', linestyle='--', marker='x')
plt.xlabel("Cap parameter m (max payoff = 2^m)")
plt.ylabel("Expected payoff")
plt.title("St. Petersburg: empirical expected payoff vs theoretical cap")
plt.legend()
plt.tight_layout()
plt.show()

# -----------------------------
# 图 2: log2(payoff) 的直方图
# -----------------------------
m_show = 30
payoffs_show = 2 ** np.minimum(N, m_show)
log2_payoffs = np.log2(payoffs_show)

plt.figure(figsize=(9,4))
plt.hist(log2_payoffs, bins=np.arange(0.5, log2_payoffs.max()+1.5, 1),
         density=True, alpha=0.6, edgecolor='black')
plt.xlabel("log2(payoff) = n (so payoff = 2^n)")
plt.ylabel("Relative frequency ~ P(N=n)")
plt.title(f"Distribution of log2(payoff) (m={m_show}, trials={trials})")
plt.tight_layout()
plt.show()

# -----------------------------
# 图 3: 稀有大 N 的贡献
# -----------------------------
m_large = 40
payoffs_large = 2 ** np.minimum(N, m_large)
total_mean = payoffs_large.mean()

ks = np.arange(1, 21)
tail_contrib = []
for k in ks:
    mask = (N >= k)
    # 占总体均值的比例
    tail_contrib.append(payoffs_large[mask].sum() / (trials * total_mean))

plt.figure(figsize=(9,4))
plt.plot(ks, tail_contrib, marker='o')
plt.xlabel("k (consider only outcomes with N >= k)")
plt.ylabel("Fraction of total mean")
plt.title("Fraction of empirical mean from rare large outcomes (m=40 cap)")
plt.grid(True)
plt.tight_layout()
plt.show()
